/* Copyright (C) 1997-2005, Antony L Hosking.
 * All rights reserved.  */
package RegAlloc;

import java.util.*;

import Translate.Temp;

public class Color {
	private LinkedHashSet<Node> simplifyWorkList = new LinkedHashSet<Node>();
	private LinkedHashSet<Node> spillWorkList = new LinkedHashSet<Node>();
	private LinkedHashSet<Node> freezeWorkList = new LinkedHashSet<Node>();
	private LinkedHashSet<Node> spilledNodes = new LinkedHashSet<Node>();
	private LinkedHashSet<Node> coalescedNodes = new LinkedHashSet<Node>();
	private Stack<Node> selectStack = new Stack<Node>();
	
	// Move Sets
	private LinkedHashSet<Move> coalescedMoves = new LinkedHashSet<Move>();
	private LinkedHashSet<Move> constrainedMoves = new LinkedHashSet<Move>();
	private LinkedHashSet<Move> frozenMoves = new LinkedHashSet<Move>();
	private LinkedHashSet<Move> worklistMoves = new LinkedHashSet<Move>();
	private LinkedHashSet<Move> activeMoves = new LinkedHashSet<Move>();
	
	private LinkedHashSet<Temp> precolored = new LinkedHashSet<Temp>();
	
    public Temp[] spills() {
        Temp[] spills = null;
        
        if(spilledNodes.isEmpty())
        	return spills;
        
        spills = new Temp[spilledNodes.size()];
        int i = 0;
        for(Node n : spilledNodes) {
        	spills[i] = n.temp;
        	i++;
        }
        
        return spills;
    }
    
    private Set<Node> Adjacent(Node node) {
    	Set<Node> adjNodes = new LinkedHashSet<Node>(node.adj());
    	// adjList[node] \ (selectStack union coalescedNodes)
    	adjNodes.removeAll(selectStack);
    	adjNodes.removeAll(coalescedNodes);
    	return adjNodes;
    }
    
    private Set<Move> NodeMoves(Node node) {
    	Set<Move> moves = new LinkedHashSet<Move>();
    	// node.moveList intersect (activeMoves union worklistMoves)
    	for(Move move : node.moveList) {
    		if(activeMoves.contains(move) || worklistMoves.contains(move)) {
    			moves.add(move);
    		}
    	}
    	return moves;
    }
    
    private boolean MoveRelated(Node node) {
    	return (!(NodeMoves(node).isEmpty()));
    }
    
    private void MakeWorklist(InterferenceGraph ig, Translate.Frame frame) {
    	for(Node node : ig.nodes()){
    		// Don't do with precolored nodes
    		if(precolored.contains(node.temp))
    			continue;
    		
    		if(node.actualDegree >= frame.registers().length){
    			spillWorkList.add(node);
    		}
    		else if(MoveRelated(node)){
    			freezeWorkList.add(node);
    		}
    		else{
    			simplifyWorkList.add(node);
    		}
    	}
    }
    
    private void Simplify() {
    	Iterator<Node> iter = simplifyWorkList.iterator();
    	Node simpNode = null;
    	if(iter.hasNext()) {
    		simpNode = iter.next();
    	}
    	else {
    		return;
    	}
    	
    	assert(!precolored.contains(simpNode.temp));
    	
		simplifyWorkList.remove(simpNode);
    	selectStack.push(simpNode);
    	
    	//Remove node from graph and adjust its adjacencies
    	Set<Node> adjNodes = Adjacent(simpNode);
    	for(Node adjNode : adjNodes) {
    		DecrementDegree(adjNode);
    	}
    }
    
    private void DecrementDegree(Node m) {
    	int d = m.actualDegree;
		m.actualDegree--;
		if(d == precolored.size()) {
			Set<Node> adjNodes = Adjacent(m);
			EnableMoves(m);
			for(Node adjNode : adjNodes) {
				EnableMoves(adjNode);
			}
			
			spillWorkList.remove(m);
			
			if(MoveRelated(m))
				freezeWorkList.add(m);
			else
				simplifyWorkList.add(m);
		}
    }
    
    private void EnableMoves(Node node) {
    	Set<Move> moves = NodeMoves(node);
    	for(Move move : moves) {
    		if(activeMoves.contains(move)) {
    			activeMoves.remove(move);
    			worklistMoves.add(move);
    		}
    	}
    }
    
    private void Coalesce(InterferenceGraph ig) {
    	Iterator<Move> iter = worklistMoves.iterator();
    	Move m = null;
    	if(iter.hasNext()) {
    		m = iter.next();
    	}
    	else {
    		return;
    	}
    	
    	Node x = GetAlias(m.dst);
    	Node y = GetAlias(m.src);
    	Node u, v;
    	
    	if(precolored.contains(y.temp)) {
    		u = y;
    		v = x;
    	}
    	else {
    		u = x;
    		v = y;
    	}
    	
    	worklistMoves.remove(m);
    	if(u == v) {
    		coalescedMoves.add(m);
    		AddWorkList(u);
    	}
    	else if(precolored.contains(v.temp) || (u.goesTo(v) || v.goesTo(u))) {
    		constrainedMoves.add(m);
    		AddWorkList(u);
    		AddWorkList(v);
    	}
    	else if((precolored.contains(u.temp) && TestOK(v, u)) || 
    			((!precolored.contains(u.temp)) && Conservative(u, v))) {
    		coalescedMoves.add(m);
    		Combine(ig, u, v);
    		AddWorkList(u);
    	}
    	else {
    		activeMoves.add(m);
    	}
    }
    
    private void AddWorkList(Node u) {
    	if((!precolored.contains(u.temp)) && (!MoveRelated(u)) && (u.actualDegree < precolored.size())) {
    		freezeWorkList.remove(u);
    		simplifyWorkList.add(u);
    	}
    }
    
    private boolean TestOK(Node v, Node u) {
    	Set<Node> adjNodes = Adjacent(v);
    	for(Node adjNode : adjNodes) {
    		if(!  ((adjNode.actualDegree < precolored.size()) || 
    			(precolored.contains(adjNode.temp)) ||
    			(adjNode.goesTo(u) || u.goesTo(adjNode)))  ) {
    			return false;
    		}
    	}
    	
    	return true;
    }
    
    private boolean Conservative(Node u, Node v) {
    	Set<Node> adjNodes = Adjacent(u);
    	adjNodes.addAll(Adjacent(v));
    	
    	int k = 0;
    	for(Node adjNode : adjNodes) {
    		if(adjNode.actualDegree >= precolored.size())
    			k++;
    	}
    	
    	return (k < precolored.size());
    }
    
    private Node GetAlias(Node n) {
    	if(n == null)
    		throw new Error("Null Alias");
    	
    	if(coalescedNodes.contains(n))
    		return GetAlias(n.alias);
    	
    	return n;
    }
    
    private void Combine(InterferenceGraph ig, Node u, Node v) {
    	if(freezeWorkList.contains(v))
    		freezeWorkList.remove(v);
    	else
    		spillWorkList.remove(v);
    	
    	coalescedNodes.add(v);
    	v.alias = u;
    	u.moveList.addAll(v.moveList);
    	
    	Set<Node> adjNodes = Adjacent(v);
    	for(Node adjNode : adjNodes) {
    		AddEdge(ig, adjNode, u);
    		DecrementDegree(adjNode);
    	}
    	
    	if((u.actualDegree >= precolored.size()) && freezeWorkList.contains(u)) {
    		freezeWorkList.remove(u);
    		spillWorkList.add(u);
    	}
    }
    
    private void AddEdge(InterferenceGraph ig, Node u, Node v) {
		if(u != v) {
			// Only Add once
			if((!u.goesTo(v)) && (!v.goesTo(u)))
				ig.addEdge(u, v);
			
			if(!precolored.contains(u.temp))
				u.actualDegree++;
			
			if(!precolored.contains(v.temp))
				v.actualDegree++;
		}
    }
    
    private void Freeze() {
    	Iterator<Node> iter = freezeWorkList.iterator();
    	Node freezeNode = null;
    	if(iter.hasNext()) {
    		freezeNode = iter.next();
    	}
    	else {
    		return;
    	}
    	
    	assert(!precolored.contains(freezeNode.temp));
    	
    	freezeWorkList.remove(freezeNode);
    	simplifyWorkList.add(freezeNode);
    	FreezeMoves(freezeNode);
    }
    
    private void FreezeMoves(Node u) {
    	for(Move m : u.moveList) {
    		Node v = m.dst;
    		if(v == u)
    			v = m.src;
    		
    		if(activeMoves.contains(m))
    			activeMoves.remove(m);
    		else
    			worklistMoves.remove(m);
    		
    		frozenMoves.add(m);
    		if(NodeMoves(v).isEmpty() && (v.actualDegree < precolored.size())) {
    			freezeWorkList.remove(v);
    			simplifyWorkList.add(v);
    		}
    	}
    }
    
    private void SelectSpill() {
    	// Select the node with minimum cost from spillWorkList
    	double spillCost = Double.MAX_VALUE;
    	Iterator<Node> iter = spillWorkList.iterator();
    	Node spillNode = null, currentNode = null;
    	while(iter.hasNext()) {
    		currentNode = iter.next();
    		assert(!precolored.contains(currentNode.temp));
    		double currentSpillCost = currentNode.spillCost / (double)currentNode.actualDegree;
    		if(currentSpillCost < spillCost) {
    			spillCost = currentSpillCost;
    			spillNode = currentNode;
    		}
    	}
    	
    	if(spillNode == null)
    		return;
    	
    	spillWorkList.remove(spillNode);
    	simplifyWorkList.add(spillNode);
    	FreezeMoves(spillNode);
    }
    
    private void AssignColors() {
    	LinkedHashSet<Node> coloredNodes = new LinkedHashSet<Node>();
    	while(!selectStack.isEmpty()) {
    		Node toColor = selectStack.pop();
    		Set<Temp> okColors = new LinkedHashSet<Temp>(precolored);
    		List<Node> adjNodes = toColor.adj();
    		for(Node adjNode : adjNodes) {
    			Node adjAlias = GetAlias(adjNode);
    			if(coloredNodes.contains(adjAlias) || okColors.contains(adjAlias.temp))
    				okColors.remove(adjAlias.color);
    		}
    		if(okColors.isEmpty()) {
    			spilledNodes.add(toColor);
    		}
    		else {
    			coloredNodes.add(toColor);
    			Iterator<Temp> iter = okColors.iterator();
    			toColor.color = iter.next();
    		}
    	}
    	for(Node n : coalescedNodes) {
    		n.color = GetAlias(n).color;
    	}
    }

    public Color(InterferenceGraph ig, Translate.Frame frame) {
	// TODO: Color each node of the interference graph ig with a register,
	// respecting the interference edges.  The colors are the
	// frame.registers().
        // Any actual spills must be returned by the callback to spills().
    	
    	// Copy working move list from ig
    	worklistMoves = new LinkedHashSet<Move>(ig.moves());
    	precolored = new LinkedHashSet<Temp>(Arrays.asList(frame.registers()));
    	
    	MakeWorklist(ig, frame);
    	
    	while((!simplifyWorkList.isEmpty()) || (!worklistMoves.isEmpty()) ||
    			(!freezeWorkList.isEmpty()) || (!spillWorkList.isEmpty())){
    		
    		if(!simplifyWorkList.isEmpty()) {
    			Simplify();
    			continue;
    		}
    		
    		if(!worklistMoves.isEmpty()) {
    			Coalesce(ig);
    			continue;
    		}
    		
    		if(!freezeWorkList.isEmpty()) {
    			Freeze();
    			continue;
    		}
    		
    		if(!spillWorkList.isEmpty()) {
    			SelectSpill();
    			continue;
    		}
    	}
    	
    	AssignColors();
    }
}
